from pathlib import Path
from PIL import Image
from PIL.ExifTags import TAGS, GPSTAGS
from typing import NamedTuple, Optional, List, Dict, Tuple
import concurrent.futures
import time
import csv
import json


class GeoCoordinate(NamedTuple):
    """地理坐标数据结构"""
    longitude: float
    latitude: float
    altitude: float
    timestamp: Optional[float] = None  # 新增：照片拍摄时间


def extract_gps_data(image_path: Path) -> Optional[Tuple[Path, GeoCoordinate]]:
    """从单张图片中提取GPS坐标，返回图片路径和坐标元组"""
    try:
        with Image.open(image_path) as img:
            # 获取EXIF数据
            exif_data = img.getexif()

            if not exif_data:
                return (image_path, None)

            # 直接获取GPSInfo IFD
            gps_info = {}
            if 0x8825 in exif_data:
                gps_ifd = exif_data.get_ifd(0x8825)
                for key, val in gps_ifd.items():
                    tag_name = GPSTAGS.get(key, key)
                    gps_info[tag_name] = val

            # 检查必要数据是否存在
            required_keys = ['GPSLatitude', 'GPSLongitude', 'GPSLatitudeRef', 'GPSLongitudeRef']
            if not all(k in gps_info for k in required_keys):
                return (image_path, None)

            # 度分秒转十进制函数
            def dms_to_decimal(dms, ref):
                def to_float(val):
                    if isinstance(val, tuple):
                        return float(val[0]) / float(val[1])
                    return float(val)

                # 确保dms有3个元素
                dms = list(dms)
                while len(dms) < 3:
                    dms.append(0)

                degrees = to_float(dms[0])
                minutes = to_float(dms[1])
                seconds = to_float(dms[2])

                decimal = degrees + minutes / 60 + seconds / 3600
                return -decimal if ref in ['S', 'W'] else decimal

            # 解析坐标
            latitude = dms_to_decimal(
                gps_info['GPSLatitude'],
                gps_info['GPSLatitudeRef']
            )
            longitude = dms_to_decimal(
                gps_info['GPSLongitude'],
                gps_info['GPSLongitudeRef']
            )

            # 解析海拔
            altitude = 0.0
            if 'GPSAltitude' in gps_info:
                if isinstance(gps_info['GPSAltitude'], tuple):
                    altitude = float(gps_info['GPSAltitude'][0]) / float(gps_info['GPSAltitude'][1])
                else:
                    altitude = float(gps_info['GPSAltitude'])

            # 尝试获取拍摄时间
            timestamp = None
            if 0x9003 in exif_data:  # DateTimeOriginal 标签
                try:
                    # 格式: "YYYY:MM:DD HH:MM:SS"
                    dt_str = exif_data[0x9003]
                    time_struct = time.strptime(dt_str, "%Y:%m:%d %H:%M:%S")
                    timestamp = time.mktime(time_struct)
                except (ValueError, TypeError):
                    pass

            return (image_path, GeoCoordinate(
                longitude=longitude,
                latitude=latitude,
                altitude=altitude,
                timestamp=timestamp
            ))

    except Exception as e:
        print(f"处理图片 {image_path.name} 时出错: {str(e)}")
        return (image_path, None)


def batch_extract_gps(directory: Path, output_format: str = "csv") -> Dict[Path, GeoCoordinate]:
    print(f"实际扫描到的文件: {list(directory.glob('*'))}")  # 添加这行
    """
    批量处理目录中的所有图片
    :param directory: 图片目录路径
    :param output_format: 输出格式 (csv, json, console)
    :return: 包含所有结果的字典
    """
    # 支持的图片格式
    image_extensions = ['.jpg', '.jpeg', '.png', '.tiff', '.heic', '.dng']

    # 收集所有图片文件
    image_files = []
    for ext in image_extensions:
        image_files.extend(directory.glob(f'*{ext}'))
        image_files.extend(directory.glob(f'*{ext.upper()}'))

    print(f"找到 {len(image_files)} 张图片")

    # 使用线程池并行处理
    results = {}
    with concurrent.futures.ThreadPoolExecutor() as executor:
        future_to_path = {executor.submit(extract_gps_data, img): img for img in image_files}

        for future in concurrent.futures.as_completed(future_to_path):
            img_path, coord = future.result()
            results[img_path] = coord

    # 输出结果
    if output_format.lower() == "csv":
        export_to_csv(results, directory / "gps_coordinates.csv")
    elif output_format.lower() == "json":
        export_to_json(results, directory / "gps_coordinates.json")
    else:  # 控制台输出
        print_results(results)

    return results


def export_to_csv(results: Dict[Path, GeoCoordinate], output_path: Path):
    """导出结果到CSV文件"""
    with open(output_path, 'w', newline='', encoding='utf-8') as csvfile:
        fieldnames = ['文件名', '纬度', '经度', '海拔(m)', '拍摄时间']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        writer.writeheader()
        for img_path, coord in results.items():
            if coord:
                # 格式化时间戳
                time_str = ""
                if coord.timestamp:
                    time_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(coord.timestamp))

                writer.writerow({
                    '文件名': img_path.name,
                    '纬度': coord.latitude,
                    '经度': coord.longitude,
                    '海拔(m)': coord.altitude,
                    '拍摄时间': time_str
                })


def export_to_json(results: Dict[Path, GeoCoordinate], output_path: Path):
    """导出结果到JSON文件"""
    output_data = []
    for img_path, coord in results.items():
        if coord:
            # 格式化时间戳
            time_str = ""
            if coord.timestamp:
                time_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(coord.timestamp))

            output_data.append({
                'file': str(img_path.name),
                'latitude': coord.latitude,
                'longitude': coord.longitude,
                'altitude': coord.altitude,
                'timestamp': coord.timestamp,
                'formatted_time': time_str
            })

    with open(output_path, 'w', encoding='utf-8') as jsonfile:
        json.dump(output_data, jsonfile, indent=2, ensure_ascii=False)


def print_results(results: Dict[Path, GeoCoordinate]):
    """在控制台输出结果"""
    print("\nGPS坐标提取结果:")
    print("-" * 80)
    for img_path, coord in results.items():
        print(f"文件: {img_path.name}")
        if coord:
            # 格式化时间戳
            time_str = ""
            if coord.timestamp:
                time_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(coord.timestamp))
                print(f"  拍摄时间: {time_str}")

            print(f"  纬度: {coord.latitude:.6f}")
            print(f"  经度: {coord.longitude:.6f}")
            print(f"  海拔: {coord.altitude:.2f} 米")
        else:
            print("  未找到GPS数据")
        print("-" * 60)


if __name__ == "__main__":
    # 使用示例
    photo_dir = Path("test_images")  # 替换为你的照片目录

    # 批量处理并导出为CSV
    results = batch_extract_gps(
        directory=photo_dir,
        output_format="csv"  # 可选: "csv", "json", "console"
    )

    # 统计结果
    total = len(results)
    success = sum(1 for coord in results.values() if coord is not None)
    print(f"\n处理完成! 成功提取 {success}/{total} 张照片的GPS数据")